from langchain_community.document_loaders import TextLoader
import os
from langchain_community.document_loaders import (
    TextLoader, 
    UnstructuredMarkdownLoader, 
    PyMuPDFLoader
)
from langchain_core.documents import Document as LangchainDocument
from rapidocr_onnxruntime import RapidOCR

def load_document(file_path: str):
    """
    根据文件类型自动选择合适的加载器来加载文档内容。
    目前只需要加载word文档
    """
    ext = os.path.splitext(file_path)[1].lower()
    
    if ext == '.docx':
        # 使用python-docx完整解析文档内容
        try:
            from docx import Document
            doc = Document(file_path)
            full_text = []
            
            # 提取所有段落文本
            for para in doc.paragraphs:
                full_text.append(para.text)
            
            # 提取表格文本
            for table in doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        full_text.append(cell.text)
            
            # 提取页眉页脚
            for section in doc.sections:
                for header in section.header.paragraphs:
                    full_text.append(header.text)
                for footer in section.footer.paragraphs:
                    full_text.append(footer.text)
            
            # 合并文本并清理格式
            cleaned_text = '\n\n'.join(full_text).replace('\x0c', '').strip()
            print(f"成功解析文档：{file_path}，提取字符数：{len(cleaned_text)}")  # 调试输出
            return [LangchainDocument(
                page_content=cleaned_text,
                metadata={"source": file_path, "type": "word"}
            )]
        except Exception as e:
            print(f"解析Word文档失败: {str(e)}")
            return []
        
    elif ext == '.txt':
        # 读取txt文件
        loader = TextLoader(file_path, encoding="utf8")
        return loader.load()
    
    elif ext == '.pdf':
        # 读取pdf文件
        loader = PyMuPDFLoader(file_path, mode="page")
        return loader.load()
    
    elif ext == '.jpg' or ext == '.jpeg' or ext == '.png':
        # 读取图片文件并使用OCR识别
        ocr = RapidOCR()
        result, _ = ocr(file_path)
        if result:
            content = "\n".join([line[1] for line in result])
            return [LangchainDocument(page_content=content)]
        return []
    
    elif ext == '.md':
        # 读取markdown文件
        loader = UnstructuredMarkdownLoader(file_path)
        return loader.load()
    else:
        raise ValueError(f"Unsupported file type: {ext}")

def load_documents_from_folder(folder_path: str):
    """
    遍历文件夹，加载所有支持的文档文件并返回其内容。
    """
    all_docs = []
    
    # 遍历文件夹中的每一个文件
    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        
        # print("============>", file_path)
        if os.path.isfile(file_path):  # 确保是文件而不是文件夹
            try:
                docs = load_document(file_path)
                for doc in docs:
                    all_docs.append({
                        "file_name": filename,
                        "content": doc.page_content,
                        "word_count": len(doc.page_content.split())
                    })  # 将文件名、内容和字数封装到字典中
            except ValueError as e:
                print(f"无法处理文件 {filename}: {e}")
    
    return all_docs

# 指定文件夹路径
folder_path = "all_file"  # 这里是你的文件夹路径
# # 从文件夹加载所有文档
docs = load_documents_from_folder(folder_path)
# print(docs)
